#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>
#include <algorithm>
#include "multi_cavity_matrix_builder.h"

// Structure to hold element-wise comparison information
struct ElementDiff {
    int index;
    float tensor_result;
    float csr_result;
    float diff;
    float rel_diff;
    
    ElementDiff(int idx, float tensor, float csr) 
        : index(idx), tensor_result(tensor), csr_result(csr) {
        diff = std::abs(tensor - csr);
        rel_diff = (std::abs(tensor) > 1e-10) ? diff / std::abs(tensor) : diff;
    }
    
    bool operator<(const ElementDiff& other) const {
        return diff > other.diff; // Sort by largest difference first
    }
};

// Analyze and print detailed comparison between tensor and CSR results
void analyze_matrix_differences(
    const MultiCavityMatrix::MultiCavityMatrixBuilder& builder,
    const float* tensor_data,
    const CsrMatrix& csr_matrix,
    const float* x,
    int size)
{
    std::cout << "\n==== DETAILED MATRIX FORMAT COMPARISON ====\n" << std::endl;
    
    // Perform matrix-vector multiplication with tensor format
    float* tensor_result = new float[size]();
    builder.matrix_vector_multiply(tensor_data, x, tensor_result);
    
    // Perform matrix-vector multiplication with CSR format
    double* x_double = new double[size];
    for (int i = 0; i < size; i++) {
        x_double[i] = static_cast<double>(x[i]);
    }
    
    double* csr_result_double = new double[size]();
    for (int i = 0; i < csr_matrix.rows; i++) {
        for (int j = csr_matrix.row_off[i]; j < csr_matrix.row_off[i + 1]; j++) {
            csr_result_double[i] += csr_matrix.data[j] * x_double[csr_matrix.cols[j]];
        }
    }
    
    float* csr_result = new float[size];
    for (int i = 0; i < size; i++) {
        csr_result[i] = static_cast<float>(csr_result_double[i]);
    }
    
    // Collect differences
    std::vector<ElementDiff> differences;
    for (int i = 0; i < size; i++) {
        ElementDiff diff(i, tensor_result[i], csr_result[i]);
        if (diff.diff > 1e-6) { // Only track meaningful differences
            differences.push_back(diff);
        }
    }
    
    // Sort by largest difference
    std::sort(differences.begin(), differences.end());
    
    // Print summary statistics
    int num_significant_diffs = 0;
    for (const auto& diff : differences) {
        if (diff.rel_diff > 0.01) { // More than 1% difference
            num_significant_diffs++;
        }
    }
    
    std::cout << "Total elements: " << size << std::endl;
    std::cout << "Elements with differences > 1e-6: " << differences.size() << std::endl;
    std::cout << "Elements with relative differences > 1%: " << num_significant_diffs << std::endl;
    
    // Print the top 20 largest differences
    int count = std::min(20, static_cast<int>(differences.size()));
    if (count > 0) {
        std::cout << "\nTop " << count << " largest differences:" << std::endl;
        std::cout << std::setw(10) << "Index" 
                  << std::setw(15) << "Tensor Result" 
                  << std::setw(15) << "CSR Result" 
                  << std::setw(15) << "Difference"
                  << std::setw(15) << "Rel Diff (%)" << std::endl;
        
        for (int i = 0; i < count; i++) {
            const auto& diff = differences[i];
            std::cout << std::setw(10) << diff.index 
                      << std::setw(15) << diff.tensor_result 
                      << std::setw(15) << diff.csr_result 
                      << std::setw(15) << diff.diff
                      << std::setw(15) << diff.rel_diff * 100.0 << "%" << std::endl;
                      
            // Get cavity and local coordinates for this point
            int cavity_id, i_idx, j_idx, k_idx;
            builder.global_to_local_3d(diff.index, cavity_id, i_idx, j_idx, k_idx);
            
            std::cout << "    (Cavity " << cavity_id << ", pos: [" 
                      << i_idx << "," << j_idx << "," << k_idx << "])" << std::endl;
                      
            // Check if this is a boundary or interface point
            bool is_interface_point = false;
            // Need to add code to check if this is an interface point
            
            if (i_idx == 0 || j_idx == 0 || k_idx == 0 || 
                i_idx == builder.get_cavity_dimensions(cavity_id).imt - 1 || 
                j_idx == builder.get_cavity_dimensions(cavity_id).jmt - 1 || 
                k_idx == builder.get_cavity_dimensions(cavity_id).kmt - 1) {
                std::cout << "    [BOUNDARY POINT]" << std::endl;
            }
            
            if (is_interface_point) {
                std::cout << "    [INTERFACE POINT]" << std::endl;
            }
        }
    }
    
    // Analyze where the differences are occurring
    int boundary_diffs = 0;
    int interior_diffs = 0;
    int interface_diffs = 0;
    
    for (const auto& diff : differences) {
        int cavity_id, i_idx, j_idx, k_idx;
        builder.global_to_local_3d(diff.index, cavity_id, i_idx, j_idx, k_idx);
        
        const MultiCavityMatrix::Grid3D& grid = builder.get_cavity_dimensions(cavity_id);
        
        // Check if this is a boundary point
        bool is_boundary = (i_idx == 0 || j_idx == 0 || k_idx == 0 || 
                          i_idx == grid.imt - 1 || j_idx == grid.jmt - 1 || k_idx == grid.kmt - 1);
                          
        // Check if this is an interface point
        bool is_interface = builder.is_interface_point(cavity_id, i_idx, j_idx, k_idx);
        
        if (is_interface) {
            interface_diffs++;
        } else if (is_boundary) {
            boundary_diffs++;
        } else {
            interior_diffs++;
        }
    }
    
    std::cout << "\nDistribution of differences:" << std::endl;
    std::cout << "  Boundary points: " << boundary_diffs << std::endl;
    std::cout << "  Interface points: " << interface_diffs << std::endl;
    std::cout << "  Interior points: " << interior_diffs << std::endl;
    
    // Examine CSR structure
    std::cout << "\nCSR Matrix Structure Analysis:" << std::endl;
    int min_entries_per_row = size;
    int max_entries_per_row = 0;
    double avg_entries_per_row = 0.0;
    
    for (int i = 0; i < csr_matrix.rows; i++) {
        int entries = csr_matrix.row_off[i+1] - csr_matrix.row_off[i];
        min_entries_per_row = std::min(min_entries_per_row, entries);
        max_entries_per_row = std::max(max_entries_per_row, entries);
        avg_entries_per_row += entries;
    }
    avg_entries_per_row /= csr_matrix.rows;
    
    std::cout << "  Non-zeros: " << csr_matrix.data_size << std::endl;
    std::cout << "  Min entries per row: " << min_entries_per_row << std::endl;
    std::cout << "  Max entries per row: " << max_entries_per_row << std::endl;
    std::cout << "  Avg entries per row: " << avg_entries_per_row << std::endl;
    
    // Check if any row has zero entries
    int zero_entry_rows = 0;
    for (int i = 0; i < csr_matrix.rows; i++) {
        if (csr_matrix.row_off[i+1] == csr_matrix.row_off[i]) {
            zero_entry_rows++;
            std::cout << "  Row " << i << " has zero entries!" << std::endl;
            if (zero_entry_rows <= 5) { // Only print details for first 5
                int cavity_id, i_idx, j_idx, k_idx;
                builder.global_to_local_3d(i, cavity_id, i_idx, j_idx, k_idx);
                std::cout << "    (Cavity " << cavity_id << ", pos: [" 
                          << i_idx << "," << j_idx << "," << k_idx << "])" << std::endl;
            }
        }
    }
    std::cout << "  Total rows with zero entries: " << zero_entry_rows << std::endl;
    
    // Clean up
    delete[] tensor_result;
    delete[] x_double;
    delete[] csr_result_double;
    delete[] csr_result;
    
    std::cout << "\n==== END OF ANALYSIS ====\n" << std::endl;
} 